import datetime
import numpy
import time
import math
from PIL import Image

import common
import utils

class VortGate:
	def __init__(self, ang1, ang2, dist1, dist2, vel1, vel2, velGrad, avgRef, localVar, link, id):
		self.angle1 = ang1
		self.angle2 = ang2
		self.dist1 = dist1
		self.dist2 = dist2
		self.center = ((ang1 + ang2)/2.0, (dist1 + dist2)/2.0)
		self.vel1 = vel1
		self.vel2 = vel2
		self.velGrad = velGrad
		self.avgRef = avgRef
		self.localVelVar = localVar
		self.link = link
		self.id = id

class ShearSegment:
	def __init__(self, vortList):
		self.vortList = vortList
		self.avgRef = 0
		self.avgVelDiff = 0
		self.avgVelGrad = 0
		self.avgVar = 0
		self.center = (0, 0)
		
class ShearFeature:
	def __init__(self):
		# these are the four radar elevations
		self.shearSegments = [None, None, None, None]
		self.avgRef = 0
		self.avgVelDiff = 0
		self.avgVelGrad = 0
		self.avgVar = 0
		self.center = (0, 0)
		self.centerLatLon = (0, 0)
		self.score = -1
		self.targetDist = -1
		
		# the id of the vort track that this feature belongs to
		self.trackId = -1


def findShearFeatures(stormRelData, baseRefData, targetLatLon = (-1, -1), maxDist = 150, verbose = False):
	# minimum relative velocity for a pixel to be considered in the vort gate search
	minVortGateRelVel = 0.0
	# minimum difference in vel betwen pixels in a vort gate
	minVortAbsVelDiff = 20.0
	# min gradiant to be considered a vort gate
	minVelGrad = 6.0
	# max gradiant to be considered a vort gate
	maxVelGrad = 70.0
	# max number of miles between two vort gates to be considered part of a shear segment
	maxVortGateSeparation = 6.0
	# number of vort gates at each level required for a shear feature
	shearFeatureVortGateMinNum = 2
	# max distance between shear segments to be in a shear feature
	maxShearSegmentSeparation = 15.0
	# number of pixels away from the target to search for a velocity gradiant
	velGradSearchDistMin = 2
	velGradSearchDistMax = 3
	# max number of zero frames in the vel grad search area
	maxVelGradZeros = 10
	# minimum score to be counted as a shear feature
	minScore = 10.0
	# max local variance for a vort gate
	maxLocalVelVar = 1500
	# max local ref variance for a vort gate
	maxLocalRefVar = 1000
	# max average shear feature var
	maxAvgVar = 950
	# min average shear feature ref
	minAvgRef = 15
	
	# angle sweep for the vortGate search
	angleSweep = 90
	
	if len(stormRelData.keys()) != 4:
		print "need storm rel data from all angles (0.5, 1.5, 2.5, 3.5)!"
		return
	
	startTime = time.time()
	
	vortGates = []
	count = 0
	sortedElevations = sorted(stormRelData.keys())
	for elevation in sortedElevations:
		curStormRelData = stormRelData[elevation]
		curVortGates = []
		
		def computeLocalStats(angle, dist, radius = 2):
			localRefData = []
			localVelData = []
		
			for a in range(angle-radius/2, angle+radius/2+1):
				for d in range(dist-radius/2, dist+radius/2+1):
					# bounds check
					if a >= 0 and a < len(curStormRelData) and \
					   d >= 0 and d < len(curStormRelData[a]):
						localVelData.append(curStormRelData[a][d])
						
					if a >= 0 and a < len(baseRefData) and \
					   d >= 0 and d < len(baseRefData[a]):
						if baseRefData[a][d] > 0:
							localRefData.append(baseRefData[a][d])
					
			velVar = 0
			refVar = 0
			velZeros = len([x for x in localVelData if x != 0])
			avgRef = 0
			
			if len(localVelData) > 0:
				velVar = numpy.var(localVelData)
			if len(localRefData) > 0:
				avgRef = numpy.mean(localRefData)
				refVar = numpy.var(localRefData)
			
			return velVar, refVar, velZeros, avgRef


		def findLocalVelMinMax(angle, dist, radius = 4):			
			minVel = -1
			minCoords = (-1, -1)			
			maxVel = -1
			maxCoords = (-1, -1)
			
			for a in range(angle-radius/2, angle+radius/2+1):
				for d in range(dist-radius/2, dist+radius/2+1):
					# bounds check
					if a < 0 or a >= len(curStormRelData) or \
					   d < 0 or d >= len(curStormRelData[a]):
						continue
					if curStormRelData[a][d] < minVel or minVel == -1:
						minVel = curStormRelData[a][d]
						minCoords = (a, d)
					if curStormRelData[a][d] > maxVel or maxVel == -1:
						maxVel = curStormRelData[a][d]
						maxCoords = (a, d)
			return {'min':minCoords, 'max':maxCoords}
		
		targetXY = utils.latLonToXY(targetLatLon[0], targetLatLon[1])
		
		maxDistY = utils.distToY(0)-utils.distToY(maxDist)
		distMin = int(max(0, targetXY[1]-maxDistY))
		distMax = int(min(common.RadarRange, targetXY[1]+maxDistY))

		angleMin = (int(targetXY[0])-angleSweep/2)%360
		angleMax = (int(targetXY[0])+angleSweep/2)%360
		angleRange = [a%360 for a in range(angleMin, angleMin+angleSweep)]
		
		for dist in range(distMin, distMax):
			for angle in angleRange:

				# don't go out of bounds
				if angle >= len(curStormRelData) or \
				   dist >= len(curStormRelData[angle]):
					continue

				# skip blank pixels
				if baseRefData[angle][dist] == 0 or abs(curStormRelData[angle][dist]) < minVortGateRelVel:
					continue

				# get distance to this pixel from targetLatLon and limit
				if targetLatLon != (-1, -1):
					curLatLon = utils.xyToLatLon(angle, dist)
					# if we've specified a target location and a maxDist, skip all pixels farther than maxDist from the target
					if utils.latLonDist(curLatLon, targetLatLon) > maxDist and maxDist != -1:
						continue
				
				vortGateChoices = []
				
				for r in range(velGradSearchDistMin, velGradSearchDistMax):
					velMinMax = findLocalVelMinMax(angle, dist, radius = r)
					
					minCoords = velMinMax['min']
					maxCoords = velMinMax['max']
					
					if minCoords == maxCoords or \
					   minCoords == (-1, -1) or \
					   maxCoords == (-1, -1):
						continue
					
					velMin = curStormRelData[minCoords[0]][minCoords[1]]
					velMax = curStormRelData[maxCoords[0]][maxCoords[1]]
					
					# zero check
					if velMin != 0 and velMax != 0:
						
						# calculate gradient
						velGrad = (velMax - velMin)/float(math.sqrt((minCoords[0]-maxCoords[0])**2.0 + (minCoords[1]-maxCoords[1])**2.0))
						# absolute velocity difference
						absVelDiff = velMax - velMin
						
						# check velocity gradient and absolute velocity difference
						if velGrad >= minVelGrad and velGrad <= maxVelGrad and absVelDiff > minVortAbsVelDiff:
							
							# check for velocity sign difference
							if velMin*velMax < 0:
								
								velVar, refVar, velZeros, avgRef = computeLocalStats(int((minCoords[0]+maxCoords[0])/2.0), 
																					 int((minCoords[1]+maxCoords[1])/2.0), r)
								
								if velZeros <= maxVelGradZeros:
									
									if velVar < maxLocalVelVar:
										
										if refVar < maxLocalRefVar or elevation != 'stormRel0':

											vortGateChoices.append(VortGate(minCoords[0], maxCoords[0], 
																			minCoords[1], maxCoords[1],
																			velMin, velMax, 
																			velGrad,
																			avgRef,
																			velVar,
																			[], count))
				
				
				if len(vortGateChoices) > 0:
					vortGateChoices.sort(key = lambda x:x.velGrad, reverse = True)
					curVortGates.append(vortGateChoices[0])
					count += 1
					vortGateFound = True	

		vortGates.append(curVortGates)
	
	if verbose:
		print 'found %d vort gates in %f seconds...'%(count, (time.time()-startTime))
	startTime = time.time()

	# search for gates that are near each other
	for elev in vortGates:
		for vort in elev:
			vortLatLon = utils.xyToLatLon(vort.center[0], vort.center[1])
			
			for vort2 in elev:			
				if vort.center[0] == vort2.center[0] and vort.center[1] == vort2.center[1]: continue

				vort2LatLon = utils.xyToLatLon(vort2.center[0], vort2.center[1])
				vortDist = utils.latLonDist(vortLatLon, vort2LatLon)
				
				if vortDist <= maxVortGateSeparation:
					if vort2.id not in vort.link: vort.link.append(vort2.id)
					if vort.id not in vort2.link: vort2.link.append(vort.id)
	if verbose:
		print 'linked vort gates in %f seconds...'%(time.time()-startTime)
	startTime = time.time()
		
	shearSegments = []
	# build shear segments by collecting all vort gates that are linked to each other
	for elev in vortGates:
		curShearSegment = []
		inShearSegment = []
		for v in elev:
			if len(v.link) == 0 or v.id in inShearSegment:
				continue
			
			# start a new shear segment
			curShearSegment.append(ShearSegment(vortList = [v]))
			inShearSegment.append(v.id)
			
			toVisit = v.link
			visited = []
			
			# visit all the nodes linked to this one, and all sublinks
			while len(toVisit) > 0:
				curVortId = toVisit.pop()
				visited.append(curVortId)
				
				curVort = None
				# find current vortGate
				for v2 in elev:
					if v2.id == curVortId:
						curVort = v2
						# loop through all of its links and add the ones we don't have yet
						for l in v2.link:
							if l not in toVisit and l not in visited:
								toVisit.append(l)
				
				# add current gate to the shear segment
				inShearSegment.append(curVortId)
				for s in range(len(curShearSegment)):
					if v in curShearSegment[s].vortList and curVort not in curShearSegment[s].vortList:
						curShearSegment[s].vortList.append(curVort)
		shearSegments.append(curShearSegment)
	
	if verbose:
		print 'built %d shear segments in %f seconds...'%(len(shearSegments), (time.time()-startTime))
	
	# calculate shear segment properties
	for e in range(len(shearSegments)):
		for i in range(len(shearSegments[e])):
			shear = shearSegments[e][i]
			avgVelDiff = 0
			avgRef = 0
			avgVelGrad = 0
			avgVar = 0
			centerAngle = 0
			centerDist = 0
			for v in shear.vortList:
				avgVelDiff += abs(v.vel1-v.vel2)
				avgRef += abs(v.avgRef)
				avgVelGrad += v.velGrad
				avgVar += v.localVelVar
				centerAngle += v.center[0]
				centerDist += v.center[1]
			avgVelDiff /= len(shear.vortList)
			avgRef /= len(shear.vortList)
			avgVelGrad /= len(shear.vortList)
			avgVar /= len(shear.vortList)
			center = (float(centerAngle)/len(shear.vortList), float(centerDist)/len(shear.vortList))
			shearSegments[e][i].avgRef = avgRef
			shearSegments[e][i].avgVelDiff = avgVelDiff
			shearSegments[e][i].avgVelGrad = avgVelGrad
			shearSegments[e][i].avgVar = avgVar
			shearSegments[e][i].center = center
	
	if verbose:
		print 'final shear segment count: %d'%len(shearSegments)
		
	shear2DFeatures	= []
	
	# link shear segments that have centers close to each other
	inFeature = []
	# loop through each shear segment in each elevation
	for e in range(len(shearSegments)):
		for s in range(len(shearSegments[e])):
			
			shearFeature = ShearFeature()
			shearSLatLon = utils.xyToLatLon(shearSegments[e][s].center[0], shearSegments[e][s].center[1])
			
			# loop through all of the other shear segments at all other elevations
			for e2 in range(len(shearSegments)):
				if e == e2: continue
				for s2 in range(len(shearSegments[e2])):
					
					shearS2LatLon = utils.xyToLatLon(shearSegments[e2][s2].center[0], shearSegments[e2][s2].center[1])
					shearDist = utils.latLonDist(shearSLatLon, shearS2LatLon)
					
					if shearDist < maxShearSegmentSeparation:
						if ((e, s) not in inFeature) and ((e2, s2) not in inFeature):
							shearFeature.shearSegments[e] = shearSegments[e][s]
							shearFeature.shearSegments[e2] = shearSegments[e2][s2]
							inFeature.append((e, s))
							inFeature.append((e2, s2))
						
			if not (shearFeature.shearSegments[0] == None and shearFeature.shearSegments[1] == None and \
				    shearFeature.shearSegments[2] == None and shearFeature.shearSegments[3] == None):
				shear2DFeatures.append(shearFeature)
	
	# make sure each shear segment is valid
	for f in shear2DFeatures[:]:
		vortGateSum = 0
		for s in f.shearSegments:
			if s != None:
				vortGateSum += len(s.vortList)
		if vortGateSum < shearFeatureVortGateMinNum:
			shear2DFeatures.remove(f)
			if verbose:
				center = (0, 0)
				for seg in f.shearSegments:
					if seg != None:
						center = utils.xyToLatLon(seg.center[0], seg.center[1])
				print 'removed feature at (%.2f, %.2f): not enough vort gates'%(center[0], center[1])

		numShearSegments = 0
		shearSegmentStatus = []
		for s in f.shearSegments:
			if s != None:
				numShearSegments += 1
				shearSegmentStatus.append(1)
			else:
				shearSegmentStatus.append(0)

		if numShearSegments < 3 and f.shearSegments[0] == None:
			shear2DFeatures.remove(f)
          
			# calculate segment center
			if verbose:
				center = (0, 0)
				for seg in f.shearSegments:
					if seg != None:
						center = utils.xyToLatLon(seg.center[0], seg.center[1])
				print 'removed feature at (%.2f, %.2f): not enough shear segments ([%d, %d, %d, %d])'%(center[0], center[1], shearSegmentStatus[0],
																										shearSegmentStatus[1], shearSegmentStatus[2],
																										shearSegmentStatus[3])
	
	# calculate 2d feature average properties
	for f in range(len(shear2DFeatures)):
		avgVelDiff = 0
		avgVelGrad = 0
		avgRef = 0
		avgVar = 0
		center = (0, 0)
		segmentNum = 0
		for segment in shear2DFeatures[f].shearSegments:
			if segment != None:
				avgVelDiff += segment.avgVelDiff
				avgVelGrad += segment.avgVelGrad
				avgRef += segment.avgRef
				avgVar += segment.avgVar
				segmentNum += 1
		avgVelDiff /= segmentNum
		avgVelGrad /= segmentNum
		avgRef /= segmentNum
		avgVar /= segmentNum
		center = shear2DFeatures[f].shearSegments[0].center
	
		shear2DFeatures[f].avgVelDiff = avgVelDiff
		shear2DFeatures[f].avgVelGrad = avgVelGrad
		shear2DFeatures[f].avgRef = avgRef
		shear2DFeatures[f].avgVar = avgVar
		shear2DFeatures[f].center = center
		shear2DFeatures[f].centerLatLon = utils.xyToLatLon(center[0], center[1])
		shear2DFeatures[f].targetDist = utils.latLonDist(targetLatLon, utils.xyToLatLon(center[0], center[1]))
		shear2DFeatures[f].score = avgRef*abs(avgVelGrad)*(1.0-shear2DFeatures[f].targetDist/maxDist)

	# remove features with low scores
	for f in shear2DFeatures[:]:
		if f.score <= minScore:
			shear2DFeatures.remove(f)
			
			if verbose:
				center = (0, 0)
				for seg in f.shearSegments:
					if seg != None:
						center = utils.xyToLatLon(seg.center[0], seg.center[1])
				print 'removed feature at (%.2f, %.2f): score too low (%.2f)'%(center[0], center[1], f.score)
				
		elif f.avgVar > maxAvgVar:
			shear2DFeatures.remove(f)
			
			if verbose:
				center = (0, 0)
				for seg in f.shearSegments:
					if seg != None:
						center = utils.xyToLatLon(seg.center[0], seg.center[1])
				print 'removed feature at (%.2f, %.2f): avgVar too high (%.2f)'%(center[0], center[1], f.avgVar)
		
		elif f.avgRef <= minAvgRef:
			shear2DFeatures.remove(f)
			
			if verbose:
				center = (0, 0)
				for seg in f.shearSegments:
					if seg != None:
						center = utils.xyToLatLon(seg.center[0], seg.center[1])
				print 'removed feature at (%.2f, %.2f): avgRef too low (%.2f)'%(center[0], center[1], f.avgRef)
		
	
	if verbose:
		print 'build %d shear features'%len(shear2DFeatures)
		print 'done with analysis'
	
	shear2DFeatures.sort(key = lambda x:x.score, reverse = True)
	return shear2DFeatures





def findShearFeaturesTest():
	import warningExtractor
	import radarExtractor
	
	print "-------- findShearFeatures() test ----------"
	
	print 'loading warnings...'
	warnings = warningExtractor.loadTorWarnings([2010])

	maxDist = 80
	minDist = 5
	
	# the warnings between 20 and 80 nm from the radar station
	distValidWarnings = []

	for warn in warnings:
		dist = utils.latLonDist(warn["countyLatLon"], common.KTLXLocation)
		if dist <= maxDist and dist >= minDist:
			distValidWarnings.append(warn)

	warn = distValidWarnings[74]
	
	print 'tracked couplets for: \nwarning for %s, %s'%(warn["countyName"], str(warn["iTime"]))
	
	# find x/y coords of warning lat/lon
	countyLatLon = warn["countyLatLon"]
	countyXYCoords = utils.latLonToXY(countyLatLon[0], countyLatLon[1])
	
	bTrack = 2
	fTrack = -2
	stormRel0 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel0')
	stormRel1 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel1')
	stormRel2 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel2')
	stormRel3 = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='stormRel3')
	baseRef = radarExtractor.getClosestRadarData(warn["iTime"], bNum = bTrack, fNum = fTrack, product='baseRef0')

	# convert the chosen radar files to NetCDF
	stormRel0Nc = radarExtractor.radarToNetCDF(stormRel0)
	stormRel1Nc = radarExtractor.radarToNetCDF(stormRel1)
	stormRel2Nc = radarExtractor.radarToNetCDF(stormRel2)
	stormRel3Nc = radarExtractor.radarToNetCDF(stormRel3)
	baseRefNc = radarExtractor.radarToNetCDF(baseRef)

	if len(stormRel0Nc) == 0 or len(stormRel1Nc) == 0 or \
	   len(stormRel2Nc) == 0 or len(stormRel3Nc) == 0 or len(baseRefNc) == 0: 
		raise Exception("vortTracker()", "Failed to gather NC file data")

	# extract the radar data from the NetCDF files
	baseRefData = []
	stormRel0Data = []
	stormRel1Data = []
	stormRel2Data = []
	stormRel3Data = []
	
	if not (len(baseRefNc) == len(stormRel0Nc) == len(stormRel1Nc) == \
			len(stormRel2Nc) == len(stormRel3Nc)):
				raise Exception("vortTester", "NC file lists are not the same length")
	
	for i in range(len(baseRefNc)):
		data, prefix = radarExtractor.getRadarData(baseRefNc[i], flat = False, flip = True)
		baseRefData.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel0Nc[i], flat = False, flip = True)
		stormRel0Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel1Nc[i], flat = False, flip = True)
		stormRel1Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel2Nc[i], flat = False, flip = True)
		stormRel2Data.append(data)
		
		data, prefix = radarExtractor.getRadarData(stormRel3Nc[i], flat = False, flip = True)
		stormRel3Data.append(data)

	shearFeatures = []
	
	for i in range(len(baseRefData)):
		print 'finding shear features for image %d'%i
		curShearFeatures = findShearFeatures({'stormRel0':stormRel0Data[i], 'stormRel1':stormRel1Data[i], 
										   'stormRel2':stormRel2Data[i], 'stormRel3':stormRel3Data[i]}, 
										   baseRefData[i], targetLatLon = countyLatLon, maxDist = 30, verbose=True)
		print 'found %d features'%len(curShearFeatures)
		for feature in curShearFeatures:
			latLonCenter = utils.xyToLatLon(feature.center[0], feature.center[1])
			print '(%f, %f);'%(latLonCenter[0], latLonCenter[1])
			print '\tgrad = %f\n\tref = %f\n\tvar = %f\n\tscore = %f'%(feature.avgVelGrad, feature.avgRef, feature.avgVar, feature.score)
	

